package com.loren.demo07.simple.handler;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.lang.NonNull;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Component
public class CustomSimpleMessageHandler extends TextWebSocketHandler {

    private static final Map<String, WebSocketSession> ONLINE_SESSION = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(@NonNull WebSocketSession session) {
        ONLINE_SESSION.put(session.getId(), session);
        log.info("connection join, sessionId = {}, onlineCount = {}", session.getId(), ONLINE_SESSION.size());
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
        String payload = message.getPayload();
        log.info("message received, from sessionId = {}, payload = {}", session.getId(), payload);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, @NonNull CloseStatus status) {
        ONLINE_SESSION.remove(session.getId());
        log.info("connection close, sessionId = {}, onlineCount = {}", session.getId(), ONLINE_SESSION.size());
    }

    @Override
    public void handleTransportError(WebSocketSession session, @NonNull Throwable exception) throws Exception {
        if (session.isOpen()) {
            session.close();
        }
        ONLINE_SESSION.remove(session.getId());
        log.error("an error has occured", exception);
    }

    /**
     * 发送给所有客户端
     */
    public void sendTextMessage(String message) throws IOException {
        this.sendTextMessage(null, message);
    }

    /**
     * 发送给指定客户端
     */
    public void sendTextMessage(String sessionId, String message) throws IOException {
        if (StringUtils.isBlank(sessionId) || !ONLINE_SESSION.containsKey(sessionId)) {
            for (Map.Entry<String, WebSocketSession> entry : ONLINE_SESSION.entrySet()) {
                this.sessionSendTextMessage(entry.getValue(), message);
            }
        } else {
            this.sessionSendTextMessage(ONLINE_SESSION.get(sessionId), message);
        }
    }

    private void sessionSendTextMessage(WebSocketSession session, String message) throws IOException {
        session.sendMessage(new TextMessage(message.getBytes()));
        log.info("send a message to sessionId = {}, message = {}", session.getId(), message);
    }
}